# encoding=utf-8
import threading

import torch
from seg_common.PredictService import PredictFactory
from seg_common import CommonAnnotation
from seg_system.segmentation.service.csnet.networks.csnet import CSNet
from PIL import Image
from torchvision import transforms
from seg_system import ApplicationConfig
import numpy as np
import cv2 as cv
import os


class CSNetService(PredictFactory):
    _instance_lock = threading.Lock()

    @CommonAnnotation.override()
    def before(self, obj, **kwargs):
        mode = kwargs.get('mode', ApplicationConfig.SystemConfig.MODE)

        if mode == ApplicationConfig.SystemConfig.MODE_FOLDER:
            return obj
        elif mode == ApplicationConfig.SystemConfig.MODE_PICTURE:
            image = Image.open(obj)
            gray = image.convert('L')
            trans = self.transform(gray)
            trans = trans.unsqueeze(0)
            trans = trans.to(ApplicationConfig.SystemConfig.DEVICE)
            return trans

    @CommonAnnotation.override()
    def forward(self, obj, **kwargs):
        with torch.no_grad():
            output = self.model(obj)
            return output

    @CommonAnnotation.override()
    def after(self, obj, **kwargs):
        o = obj.detach().cpu()
        l = []
        for each_o in o:
            each_o = self.return_prediction(each_o)
            l.append(each_o)
        return l

    def return_prediction(self, pred):
        # for MSELoss()
        mask = pred.permute(1, 2, 0).contiguous()
        mask = mask.squeeze_(0).squeeze_(-1)
        mask = mask.data.cpu().numpy() * 255
        mask[mask < 127] = 0
        mask[mask >= 127] = 255
        return mask

    def to_list(self, obj):
        if not isinstance(obj, list):
            return [obj]
        return obj

    def save(self, obj, path, name, **kwargs):
        obj = self.to_list(obj)
        name = self.to_list(name)

        each_from = kwargs.get('each_from', None)
        tmp_path = kwargs.get('tmp_path', None)
        assert each_from is not None
        assert tmp_path is not None

        alpha = kwargs.get('alpha', 0.5)
        beta = 1.0 - alpha
        color = kwargs.get('color', np.array([0, 212, 255]))  # yellow bgr

        for each_o, each_n in zip(obj, name):
            file = os.path.join(path, each_n)
            tmp_file = os.path.join(tmp_path, each_n)

            ori_file = cv.imread(os.path.join(each_from, each_n))
            no_zero_index = np.nonzero(each_o)
            ori_file[no_zero_index[0], no_zero_index[1], :] = \
                ori_file[no_zero_index[0], no_zero_index[1], :] * beta + alpha * color

            cv.imwrite(tmp_file, ori_file)
            cv.imwrite(file, each_o)
            assert os.path.exists(file)

    def __init__(self, model_path: str):
        self.model = CSNet(1, 1)
        # 2022.9.14根据继承体系修改而来，可能浪费显存
        model_path = os.path.join(model_path, ApplicationConfig.NetConfig.SEGMENTATION_DICT[
            ApplicationConfig.NetConfig.SEGMENTATION_USE_CSNET])

        if model_path != "":
            model_d = torch.load(model_path)
            self.model.load_state_dict(model_d)
        self.model.to(ApplicationConfig.SystemConfig.DEVICE)
        self.model.eval()
        self.transform = transforms.Compose([
            transforms.Resize([384, 384]),
            transforms.ToTensor()
        ])

    def __new__(cls, *args, **kwargs):
        """实际上就是一个标准流程工具类，所以声明成一个单例
            修改为可继承的单例
        """
        if not hasattr(cls, 'instance_dict'):
            CSNetService.instance_dict = {}

        if str(cls) not in CSNetService.instance_dict.keys():
            with CSNetService._instance_lock:
                _instance = super().__new__(cls)
                CSNetService.instance_dict[str(cls)] = _instance
        return CSNetService.instance_dict[str(cls)]
